"""
可视化分割掩码图像
显示原始图像、分割掩码和叠加效果
"""

import os
import numpy as np
import matplotlib.pyplot as plt
import tifffile
from pathlib import Path
import random

def visualize_masks(raw_data_path, mask_data_path, output_dir="mask_visualizations"):
    """
    可视化分割掩码
    
    Args:
        raw_data_path: 原始图像路径
        mask_data_path: 掩码图像路径  
        output_dir: 输出目录
    """
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)
    
    # 获取所有图像文件
    raw_files = sorted([f for f in os.listdir(raw_data_path) if f.endswith('.tif')])
    mask_files = sorted([f for f in os.listdir(mask_data_path) if f.endswith('.tif')])
    
    print(f"找到 {len(raw_files)} 个原始图像文件")
    print(f"找到 {len(mask_files)} 个掩码文件")
    
    # 随机选择几个时间点进行可视化
    sample_indices = random.sample(range(min(len(raw_files), len(mask_files))), min(5, len(raw_files)))
    
    for idx in sample_indices:
        if idx >= len(raw_files) or idx >= len(mask_files):
            continue
            
        raw_file = raw_files[idx]
        mask_file = mask_files[idx]
        
        # 读取图像
        raw_img = tifffile.imread(os.path.join(raw_data_path, raw_file))
        mask_img = tifffile.imread(os.path.join(mask_data_path, mask_file))
        
        # 创建可视化
        fig, axes = plt.subplots(1, 3, figsize=(15, 5))
        
        # 原始图像
        axes[0].imshow(raw_img, cmap='gray')
        axes[0].set_title(f'原始图像 - {raw_file}')
        axes[0].axis('off')
        
        # 分割掩码
        axes[1].imshow(mask_img, cmap='tab20')
        axes[1].set_title(f'分割掩码 - {mask_file}')
        axes[1].axis('off')
        
        # 叠加效果
        overlay = raw_img.copy()
        if len(overlay.shape) == 2:
            overlay = np.stack([overlay] * 3, axis=-1)
        
        # 为每个细胞ID分配随机颜色
        unique_ids = np.unique(mask_img)
        unique_ids = unique_ids[unique_ids != 0]  # 排除背景
        
        for cell_id in unique_ids:
            mask = mask_img == cell_id
            color = np.random.rand(3)
            for c in range(3):
                overlay[mask, c] = overlay[mask, c] * 0.7 + color[c] * 0.3
        
        axes[2].imshow(overlay)
        axes[2].set_title(f'叠加效果 - {raw_file}')
        axes[2].axis('off')
        
        plt.tight_layout()
        plt.savefig(os.path.join(output_dir, f'visualization_{raw_file.replace(".tif", ".png")}'), 
                   dpi=150, bbox_inches='tight')
        plt.close()
        
        print(f"已保存可视化: {raw_file}")
    
    print(f"\n可视化结果保存在: {output_dir}")

def analyze_tracking_results(mask_data_path):
    """
    分析追踪结果
    
    Args:
        mask_data_path: 掩码数据路径
    """
    mask_files = sorted([f for f in os.listdir(mask_data_path) if f.endswith('.tif')])
    
    print(f"\n=== 追踪结果分析 ===")
    print(f"总时间帧数: {len(mask_files)}")
    
    # 统计每个时间点的细胞数量
    cell_counts = []
    for mask_file in mask_files:
        mask_img = tifffile.imread(os.path.join(mask_data_path, mask_file))
        unique_ids = np.unique(mask_img)
        cell_count = len(unique_ids[unique_ids != 0])  # 排除背景
        cell_counts.append(cell_count)
    
    print(f"平均细胞数量: {np.mean(cell_counts):.1f}")
    print(f"最大细胞数量: {np.max(cell_counts)}")
    print(f"最小细胞数量: {np.min(cell_counts)}")
    
    # 绘制细胞数量随时间的变化
    plt.figure(figsize=(10, 6))
    plt.plot(cell_counts, 'b-', linewidth=2)
    plt.xlabel('时间帧')
    plt.ylabel('细胞数量')
    plt.title('细胞数量随时间的变化')
    plt.grid(True, alpha=0.3)
    plt.savefig('cell_count_over_time.png', dpi=150, bbox_inches='tight')
    plt.close()
    
    print("细胞数量变化图已保存为: cell_count_over_time.png")

if __name__ == "__main__":
    # 设置中文字体
    plt.rcParams['font.sans-serif'] = ['SimHei']
    plt.rcParams['axes.unicode_minus'] = False
    
    # 数据路径
    raw_data_path = "ctc_raw_data/challenge/3t3/08"
    mask_data_path = "ctc_raw_data/challenge/3t3/08_RES"
    
    if os.path.exists(raw_data_path) and os.path.exists(mask_data_path):
        print("开始可视化分割掩码...")
        visualize_masks(raw_data_path, mask_data_path)
        analyze_tracking_results(mask_data_path)
    else:
        print("错误: 数据路径不存在")
        print(f"原始图像路径: {raw_data_path}")
        print(f"掩码图像路径: {mask_data_path}") 